package xh.ketama;

import java.nio.charset.Charset;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;

/**
 * Created by xh on 2016/3/7.
 */
public class KetamaRouter {

    private List<Node> nodes;

    private SortedMap<Long, Node> circle;

    private static ThreadLocal<MessageDigest> md5Local = new ThreadLocal<MessageDigest>();

    private int avgFactor;

    public KetamaRouter(List<Node> nodes) {
        this(nodes, 256);
    }

    public KetamaRouter(List<Node> nodes, int avgFactor) {
        this.nodes = nodes;
        this.avgFactor = avgFactor;
        this.circle = new TreeMap<Long, Node>();
        initCircle();
    }

    public Node hitNode(String key) {
        byte[] bKey = md5(key);

        long t = (long) (bKey[3] & 0xFF) << 24
                | (long) (bKey[2] & 0xFF) << 16
                | (long) (bKey[1] & 0xFF) << 8
                | (long) (bKey[0] & 0xFF);

        SortedMap<Long, Node> tailMap = circle.tailMap(t);
        if (tailMap.isEmpty()) {
            return circle.get(circle.firstKey());
        } else {
            return tailMap.get(tailMap.firstKey());
        }
    }

    private void initCircle() {

        float times = calWeightTimes();

        for (Node n : nodes) {
            int piece = (int) (n.weight * times);
            if (piece < 4) {
                piece = 4;
            }

            for (int i = 0; i < piece / 4; i++) {

                String keyForPieces = n.name + "-" + i;
                byte[] md5Bytes = md5(keyForPieces);

                for (int j = 0; j < 4; j++) {
                    long t = (long) (md5Bytes[j * 4 + 3] & 0xFF) << 24
                            | (long) (md5Bytes[j * 4 + 2] & 0xFF) << 16
                            | (long) (md5Bytes[j * 4 + 1] & 0xFF) << 8
                            | (long) (md5Bytes[j * 4] & 0xFF);

                    circle.put(t, n);
                }
            }
        }

        Set<Map.Entry<Long, Node>> entries = circle.entrySet();
        for (Map.Entry<Long, Node> entry : entries) {
            System.out.println(entry.getKey() + " : " + entry.getValue().name);
        }

    }

    private float calWeightTimes() {
        int sum = 0;
        for (Node n : nodes) {
            sum += n.weight;
        }

        int expectSum = nodes.size() * avgFactor;
        if (expectSum > sum) {
            return (float) expectSum / sum;
        } else {
            return 1f;
        }
    }

    /**
     * @return 16 bits
     */
    private static byte[] md5(String k) {
        MessageDigest md5 = md5Local.get();
        if (md5 == null) {
            try {
                md5 = MessageDigest.getInstance("MD5");
                md5Local.set(md5);
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException("MD5 not supported", e);
            }
        }
        md5.reset();
        md5.update(k.getBytes(Charset.forName("utf-8")));
        return md5.digest();
    }
}
